#include <QtWidgets>
#include "sm4.h"
#include "widget.h"
#include "ui_widget.h"

Widget::Widget(QWidget *parent) :
    QWidget(parent),
    ui(new Ui::Widget)
{
    ui->setupUi(this);
    k1 = new char[17];
    server = new QTcpServer(this);
    socket = new QTcpSocket(this);
    //有新连接产生
    connect(server, &QTcpServer::newConnection, this, [this]()
    {
        ui->btn_send_to_client->setEnabled(true);
        newSocket = server->nextPendingConnection();
        connect(newSocket, &QTcpSocket::readyRead, this, [this]()
        {
            QByteArray input = newSocket->readAll();
            QStringList info = QString(input).split(',');
            if (info.count() == 4 && info.first() == "--connect--")
            {
                ui->btn_send_to_server->setEnabled(true);
                ui->btn_disconnect->setEnabled(true);
                ui->btn_connect->setEnabled(false);
                Crypto server(info[2].toULong(), info[3].toULong());
                size_t y = server.get_pub_key();
                memcpy(k1, server.get_key(info[1].toULong()), 16);
                k1[16] = 0;
                newSocket->write(QString("--connect--,%1").arg(y).toUtf8());
                newSocket->waitForBytesWritten();
                return;
            }
            input = input.toHex();
            while (!input.isEmpty())
            {
                ui->recv_from_client->append(input.mid(0, 32));
                input.remove(0, 32);
            }
        });
    });
    //读取服务器发送的数据
    connect(socket, &QTcpSocket::readyRead, this, [this]()
    {
        QByteArray input = socket->readAll();
        QStringList info = QString(input).split(',');
        if (info.count() == 2 && info.first() == "--connect--")
        {
            k2 = client.get_key(info.last().toULong());
            return;
        }
        input = input.toHex();
        while (!input.isEmpty())
        {
            ui->recv_from_server->append(input.mid(0, 32));
            input.remove(0, 32);
        }
    });
    //客户端状态发生改变
    connect(socket, &QTcpSocket::stateChanged, this, [this](QAbstractSocket::SocketState socketState)
    {
        if (socketState == QAbstractSocket::UnconnectedState)
        {
            ui->btn_send_to_server->setEnabled(false);
            ui->btn_disconnect->setEnabled(false);
            ui->btn_connect->setEnabled(true);
            return;
        }
        if (socketState == QAbstractSocket::ConnectedState)
        {
            size_t* g_p = client.get_public();
            size_t x = client.get_pub_key();
            QStringList info;
            info.push_back("--connect--");
            info.push_back(QString::number(x));
            info.push_back(QString::number(g_p[0]));
            info.push_back(QString::number(g_p[1]));
            socket->write(info.join(',').toUtf8());
            socket->waitForBytesWritten();
            return;
        }
    });
    //开启服务器
    connect(ui->btn_open, &QPushButton::clicked, this, [this]()
    {
        if (!server->listen(QHostAddress(ui->ip->text()), ui->port->text().toUShort()))
        {
            QMessageBox::information(this, tr("提示"), tr("服务器开启失败!"));
            return;
        }
        ui->btn_open->setEnabled(false);
        ui->btn_close->setEnabled(true);
    });
    //关闭服务器
    connect(ui->btn_close, &QPushButton::clicked, this, [this]()
    {
        server->close();
        socket->disconnectFromHost();
        ui->btn_open->setEnabled(true);
        ui->btn_close->setEnabled(false);
        ui->btn_send_to_client->setEnabled(false);
    });
    //连接服务器
    connect(ui->btn_connect, &QPushButton::clicked, this, [this]()
    {
        socket->abort();
        ui->btn_connect->setEnabled(false);
        socket->connectToHost(ui->server_ip->text(), ui->server_port->text().toUShort());
    });
    //断开服务器
    connect(ui->btn_disconnect, &QPushButton::clicked, this, [this]()
    {
        ui->btn_send_to_client->setEnabled(false);
        ui->btn_disconnect->setEnabled(false);
        socket->disconnectFromHost();
    });
    //发送到客户端
    connect(ui->btn_send_to_client, &QPushButton::clicked, this, [this]()
    {
        //对字符进行分包发送 16字符为一包
        QByteArray input = ui->input_server->toPlainText().toUtf8();
        while (!input.isEmpty())
        {
            QByteArray data = input.mid(0, 16);
            if (data.count() < 16)
            {
                for (int i = 0; i < 16; ++i)
                {
                    if (data.count() > i)
                        continue;
                    data.push_back('\0'); //用空字符补够16
                }
            }
            newSocket->write(encrypt(data));
            newSocket->waitForBytesWritten();
            input.remove(0, 16);
        }
        ui->input_server->clear();
    });
    //发送到服务端
    connect(ui->btn_send_to_server, &QPushButton::clicked, this, [this]()
    {
        //对字符进行分包发送 16字符为一包
        QByteArray input = ui->input_client->toPlainText().toUtf8();
        while (!input.isEmpty())
        {
            QByteArray data = input.mid(0, 16);
            if (data.count() < 16)
            {
                for (int i = 0; i < 16; ++i)
                {
                    if (data.count() > i)
                        continue;
                    data.push_back('\0'); //用空字符补够16
                }
            }
            socket->write(encrypt(data));
            socket->waitForBytesWritten();
            input.remove(0, 16);
        }
        ui->input_client->clear();
    });
    //清空接收
    connect(ui->btn_clear_server, &QPushButton::clicked, this, [this]()
    {
        ui->recv_from_client->clear();
        ui->decrypt_server->clear();
    });
    connect(ui->btn_clear_client, &QPushButton::clicked, this, [this]()
    {
        ui->recv_from_server->clear();
        ui->decrypt_client->clear();
    });
    //解密
    connect(ui->btn_decrypt_sever, &QPushButton::clicked, this, [this]()
    {
        ui->decrypt_server->clear();
        for (const QString &data : ui->recv_from_client->toPlainText().split('\n'))
            ui->decrypt_server->append(QTextCodec::codecForName("utf-8")->toUnicode(decrypt(QByteArray::fromHex(data.toLocal8Bit()))));
    });
    connect(ui->btn_decrypt_client, &QPushButton::clicked, this, [this]()
    {
        ui->decrypt_client->clear();
        for (const QString &data : ui->recv_from_server->toPlainText().split('\n'))
            ui->decrypt_client->append(QTextCodec::codecForName("utf-8")->toUnicode(decrypt(QByteArray::fromHex(data.toLocal8Bit()))));
    });
}

Widget::~Widget()
{
    delete ui;
}

QByteArray Widget::encrypt(const QByteArray &data)
{
    sm4_context ctx;
    sm4_setkey_enc(&ctx, (unsigned char *)k1);
    sm4_crypt_ecb(&ctx,
                  SM4_ENCRYPT,
                  16,
                  (unsigned char *)data.data(),
                  (unsigned char *)data.data());
    return data;
}

QByteArray Widget::decrypt(const QByteArray &data)
{
    sm4_context ctx;
    sm4_setkey_dec(&ctx, (unsigned char *)k2);
    sm4_crypt_ecb(&ctx,
                  SM4_DECRYPT,
                  16,
                  (unsigned char *)data.data(),
                  (unsigned char *)data.data());
    return data;
}
